import sys
import os
# 添加 project/ 目录到 sys.path
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))

import argparse
import multiprocessing
from tictoc import TicToc
from tqdm import tqdm


def main(args):
    # 统计耗时
    cost = TicToc("生成fake_txt")
    assert os.path.exists(args.data_path)
    frames = os.listdir(args.data_path)
    frames.sort(key=lambda x: x)
    files = []
    for dir in frames:
        if dir[:2] == '__':  # '__'
            dir = os.path.join(args.data_path, dir)
            if os.path.isdir(dir):
                files.append(dir)

    process_size = len(files)
    manager = multiprocessing.Manager()
    if process_size > 1:
        pool = multiprocessing.Pool(process_size)
        counter_list = manager.list()
        for idx in range(process_size):
            pool.apply_async(main_worker, args=(files[idx], args))
        pool.close()
        pool.join()
    else:
        main_worker(files[0], args)

    print("---------------------------------------------------------")
    print("处理完成: {}".format(files))
    cost.toc()
    print("---------------------------------------------------------")


def main_worker(frames_path, args):
    camera_topics = ["camera75", "camera77", "camera80", "camera81"]
    paths_list_test = []

    for topic in camera_topics:
        path = os.path.join(frames_path, 'camera_epe', topic, 'GBuffers')
        file_generator = (os.path.join(root, file) for root, dirs, files in os.walk(os.path.abspath(path)) for file in
                          files)
        lines = path.strip().split('/')
        camera_ori_file = os.path.join(frames_path, 'camera', topic)
        camera_epe_file = os.path.join(frames_path, 'camera_epe', topic)
        for file_path in tqdm(file_generator, postfix=os.path.join(lines[-4], lines[-2]) + "/ 生成fake_txt中..."):
            # 确保文件是图片文件，这里假设你的图片文件是 .jpg 或 .png 格式
            name = os.path.basename(file_path).split(".")[0]
            paths_list_test.append(os.path.join(camera_ori_file, name + '.jpg') + ","
                                   + os.path.join(camera_ori_file, name + '.jpg') + ","
                                   + os.path.join(camera_epe_file, 'GBuffers', name + '.npz') + ","
                                   + os.path.join(camera_epe_file, 'Semantic', name + '.npz'))

    with open(os.path.join(frames_path, 'camera_epe', 'test.txt'), 'w') as file:
        for element in paths_list_test:
            file.write(element + '\n')


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--data-path', default="/home/adt/bags/work_space/datasets",
                        help='your data root for carla')
    args = parser.parse_args()
    main(args)
